#!/usr/bin/env python
# coding: utf-8

# # 优化器
# 
# 模型训练过程中，使用优化器计算梯度并更新网络参数，合适的优化器可以有效减少训练时间，提高模型性能。
# 
# 最基本的优化器是随机梯度下降算法（SGD），很多优化器在SGD的基础上进行了改进，以实现目标函数能更快速更有效地收敛到全局最优点。luojianet_ms中的`nn`模块提供了常用的优化器，如`nn.SGD`、`nn.Adam`、`nn.Momentum`等。本章主要介绍如何配置luojianet_ms提供的优化器以及如何自定义优化器。
# 
# ![learningrate.png](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r1.7/tutorials/source_zh_cn/advanced/network/images/learning_rate.png)
# 
# > luojianet_ms提供的优化器详细内容参见[优化器API](http://58.48.42.237/luojiaNet/luojiaNetapi/#优化器)。
# 
# ## 配置优化器
# 
# 使用luojianet_ms提供的优化器时，首先需要指定待优化的网络参数`params`，然后设置优化器的其他主要参数，如学习率`learning_rate`和权重衰减`weight_decay`等。
# 
# 如果要为不同网络参数单独设置选项，如对卷积参数和非卷积参数设置不同的学习率，可使用参数分组的方法来设置优化器。
# 
# ### 参数配置
# 
# 在构建优化器实例时，需要通过优化器参数`params`配置模型网络中要训练和更新的权重，常见的参数配置方法有以下两种。
# 
# - 使用`trainable_params`配置参数
# 
# 变量`Parameter`中包含了一个`requires_grad`的布尔型的类属性，表示模型中的网络参数是否需要进行更新。luojianet_ms中的`trainable_params`方法会屏蔽掉`Parameter`中`requires_grad`为False的属性，默认在网络模型训练的时候自动更新所有`Parameter`中的参数。
# 
# 在构建优化器时，可使用`trainable_params`方法来指定需要优化和更新的网络参数。
# In[1中低阶API实现深度学习]:


import numpy as np
import luojianet_ms.ops as ops
from luojianet_ms import nn, Tensor, Parameter

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.matmul = ops.MatMul()
        self.conv = nn.Conv2d(1, 6, 5, pad_mode="valid")
        self.param = Parameter(Tensor(np.array([1.0], np.float32)))

    def forward(self, x):
        x = self.conv(x)
        x = x * self.param
        out = self.matmul(x, x)
        return out

net = Net()

# 配置优化器需要更新的参数
optim = nn.Adam(params=net.trainable_params())
print(net.trainable_params())


# - 自定义筛选参数
#
# 用户也可以设定筛选条件，在使用`get_parameters`获取到网络全部参数后，通过限定参数名字等方法，自定义`filter`来决定哪些参数需要更新。
#
# 例如下面的例子，训练过程中将只对非卷积参数进行更新。

# In[2高级数据集管理]:


params_all = net.get_parameters()
no_conv_params = list(filter(lambda x: 'conv' not in x.name, params_all))

optim = nn.Adam(params=no_conv_params)
print(no_conv_params)


# ### 学习率
#
# 学习率作为机器学习及深度学习中常见的超参，对目标函数能否收敛到局部最小值及何时收敛到最小值有重要影响。学习率过大容易导致目标函数波动较大，难以收敛到最优值，太小则会导致收敛过程耗时过长。除了设置固定学习率，luojianet_ms还支持设置动态学习率，这些方法在深度学习网络中能明显提升收敛效率。
#
# #### 固定学习率
#
# 使用固定学习率时，优化器传入的`learning_rate`为浮点类型或标量Tensor。
#
# 以`nn.Momentum`为例，固定学习率为0.01，示例如下：

# In[3图像处理]:


# 设置学习率为0.01
optim = nn.Momentum(params=net.trainable_params(), learning_rate=0.01, momentum=0.9)


# #### 动态学习率
# # `luojianet_ms.nn`提供了动态学习率的模块，分为Dynamic LR函数和LearningRateSchedule类。其中Dynamic LR函数会预先生成长度为`total_step`的学习率列表，将列表传入优化器中使用，训练过程中，第i步使用第i个学习率的值作为当前step的学习率，其中`total_step`的设置值不能小于训练的总步数；LearningRateSchedule类将实例传递给优化器，优化器根据当前step计算得到当前的学习率。
#
# - Dynamic LR函数
#
# [Dynamic LR]函数目前有基于余弦衰减函数计算学习率（`nn.cosine_decay_lr`）、基于指数衰减函数计算学习率（`nn.exponential_decay_lr`）、基于逆时衰减函数计算学习率（`nn.inverse_decay_lr`）、基于自然指数衰减函数计算学习率（`nn.natural_exp_decay_lr`）、获取分段常量学习率（`nn.piecewise_constant_lr`）、基于多项式衰减函数计算学习率（`nn.polynomial_decay_lr`）和预热学习率(`nn.warmup_lr`)。
#
# 如下示例以分段常量学习率`nn.piecewise_constant_lr`为例：

# In[4自然语言]:


from luojianet_ms import nn

milestone = [1, 3, 10]
learning_rates = [0.1, 0.05, 0.01]
lr = nn.piecewise_constant_lr(milestone, learning_rates)

# 打印学习率
print(lr)

net = Net()
# 优化器设置待优化的网络参数和分段常量学习率
optim = nn.SGD(net.trainable_params(), learning_rate=lr)


# - LearningRateSchedule类
#
# [LearningRateSchedule类]目前有基于余弦衰减函数计算学习率（`nn.CosineDecayLR`）、基于指数衰减函数计算学习率（`nn.ExponentialDecayLR`）、基于逆时衰减函数计算学习率（`nn.InverseDecayLR`）、基于自然指数衰减函数计算学习率（`nn.NaturalExpDecayLR`）、基于多项式衰减函数计算学习率（`nn.PolynomialDecayLR`）和预热学习率(`nn.WarmUpLR`)。
#
# 如下示例基于指数衰减函数计算学习率`nn.ExponentialDecayLR`为例：

# In[5]:


from luojianet_ms import dtype
from luojianet_ms import Tensor

learning_rate = 0.1  # 学习率的初始值
decay_rate = 0.9     # 衰减率
decay_steps = 4      # 衰减的step数
step_per_epoch = 2

exponential_decay_lr = nn.ExponentialDecayLR(learning_rate, decay_rate, decay_steps)

for i in range(decay_steps):
    step = Tensor(i, dtype.int32)
    result = exponential_decay_lr(step)
    print(f"step{i+1}, lr:{result}")

net = Net()

# 优化器设置学习率为基于指数衰减函数计算学习率
optim = nn.Momentum(net.trainable_params(), learning_rate=exponential_decay_lr, momentum=0.9)



# ### 权重衰减
# 
# 权重衰减(weight decay)，通常也被称为L2正则化，是一种减少深度学习神经网络模型过拟合的方法。
# 
# 一般情况下，`weight_decay`取值范围为$[0, 1中低阶API实现深度学习)$，其默认值为0.0，此时不使用权重衰减策略。

# In[6]:


net = Net()
optimizer = nn.Momentum(net.trainable_params(), learning_rate=0.01,
                        momentum=0.9, weight_decay=0.9)


# 此外，luojianet_ms还支持动态weight decay。此时weight_decay是用户自定义的一个Cell，称之为weight_decay_schedule。在训练过程中，优化器调内部会用该Cell的实例，传入global_step计算当前step的weight_decay值。其中global_step是内部维护的变量，每训练一个step，global_step会自加1。注意，自定义的weight_decay_schedule的forward仅接收一个输入。如下是weight_decay在训练过程中进行指数衰减的一个示例。

# In[ ]:
from luojianet_ms.nn import Module
from luojianet_ms import ops, nn
import luojianet_ms.common.dtype as mstype

class ExponentialWeightDecay(Module):

    def __init__(self, weight_decay, decay_rate, decay_steps):
        super(ExponentialWeightDecay, self).__init__()
        self.weight_decay = weight_decay
        self.decay_rate = decay_rate
        self.decay_steps = decay_steps
        self.pow = ops.Pow()
        self.cast = ops.Cast()

    def forward(self, global_step):
        # forward只能有一个输入，训练过程中，会自动传入global step进行计算
        p = self.cast(global_step, mstype.float32) / self.decay_steps
        return self.weight_decay * self.pow(self.decay_rate, p)

net = Net()

weight_decay = ExponentialWeightDecay(weight_decay=0.0001, decay_rate=0.1, decay_steps=10000)
optimizer = nn.Momentum(net.trainable_params(), learning_rate=0.01,
                        momentum=0.9, weight_decay=weight_decay)


# ### 超参分组
# 
# 优化器也支持为不同参数单独设置选项，此时不直接传入变量，而是传入一个字典的列表，每个字典定义一个参数组别的设置值，key可以为`params`、`lr`、`weight_decay`和`grad_centralizaiton`，value为对应的设定值。
# 
# 其中`params`必须配置，其余参数可以选择配置，未配置的参数项，将采用定义优化器时设置的参数值。分组时，学习率可以使用固定学习率，也可以使用动态学习率。`weight_decay`可以使用固定值。
# 
# 如下示例分别对卷积参数和非卷积参数设置不同的学习率和权重衰减参数。

# In[7]:


net = Net()

# 卷积参数
conv_params = list(filter(lambda x: 'conv' in x.name, net.trainable_params()))
# 非卷积参数
no_conv_params = list(filter(lambda x: 'conv' not in x.name, net.trainable_params()))

# 固定学习率
fix_lr = 0.01

# 基于多项式衰减函数计算学习率
polynomial_decay_lr = nn.PolynomialDecayLR(learning_rate=0.1,      # 学习率初始值
                                           end_learning_rate=0.01, # 学习率最终值
                                           decay_steps=4,          # 衰减的step数
                                           power=0.5)              # 多项式幂

# 卷积参数使用固定学习率0.001,权重衰减为0.01
# 非卷积参数使用动态学习率，权重衰减为0.0
group_params = [{'params': conv_params, 'weight_decay': 0.01, 'lr': fix_lr},
                {'params': no_conv_params, 'lr': polynomial_decay_lr}]

optim = nn.Momentum(group_params, learning_rate=0.1, momentum=0.9, weight_decay=0.0)


# > 当前luojianet_ms除个别优化器外（例如AdaFactor，FTRL），均支持对学习率进行分组，详情参考[优化器API](https://www.luojianet_ms.cn/docs/zh-CN/r1.7/api_python/luojianet_ms.nn.html#优化器)。
# 
# ## 自定义优化器
# 
# 除使用luojianet_ms提供的优化器外，用户可以自定义优化器。
# 
# 自定义优化器时需要继承优化器基类[nn.Optimizer]，并重写`__init__`方法和`forward`方法实现参数的更新。
# 
# 如下示例实现自定义优化器Momentum（带有动量的SGD算法）：
# 
# $$ v_{t+1中低阶API实现深度学习} = v_t×u+grad \tag{1中低阶API实现深度学习} $$
# 
# $$p_{t+1中低阶API实现深度学习} = p_t - lr*v_{t+1中低阶API实现深度学习} \tag{2高级数据集管理} $$
# 
# 其中，$grad$ 、$lr$ 、$p$ 、$v$ 和 $u$ 分别表示梯度、学习率、权重参数、动量参数（Momentum）和初始速度。

# In[8]:


from luojianet_ms import Tensor, Parameter
from luojianet_ms import nn, ops
from luojianet_ms import dtype as mstype

class Momentum(nn.Optimizer):
    """定义优化器"""
    def __init__(self, params, learning_rate, momentum=0.9):
        super(Momentum, self).__init__(learning_rate, params)
        self.momentum = Parameter(Tensor(momentum, mstype.float32), name="momentum")
        self.moments = self.parameters.clone(prefix="moments", init="zeros")
        self.assign = ops.Assign()

    def forward(self, gradients):
        """forward输入为梯度，在训练中自动传入梯度gradients"""
        lr = self.get_lr()
        params = self.parameters # 待更新的权重参数

        for i in range(len(params)):
            # 更新moments值
            self.assign(self.moments[i], self.moments[i] * self.momentum + gradients[i])
            update = params[i] - self.moments[i] * lr  #带有动量的SGD算法
            self.assign(params[i], update)
        return params

net = Net()
# 设置优化器待优化的参数和学习率为0.01
opt = Momentum(net.trainable_params(), 0.01)


# `luojianet_ms.ops`也封装了优化器算子供用户自行定义优化器，如`ops.ApplyCenteredRMSProp`、 `ops.ApplyMomentum`和`ops.ApplyRMSProp`等。如下示例使用`ApplyMomentum`算子自定义优化器Momentum：

# In[9]:


class Momentum(nn.Optimizer):
    """定义优化器"""
    def __init__(self, params, learning_rate, momentum=0.9):
        super(Momentum, self).__init__(learning_rate, params)
        self.moments = self.parameters.clone(prefix="moments", init="zeros")
        self.momentum = momentum
        self.opt = ops.ApplyMomentum()

    def forward(self, gradients):
        # 待更新的权重参数
        params = self.parameters
        success = None
        for param, mom, grad in zip(params, self.moments, gradients):
            success = self.opt(param, mom, self.learning_rate, grad, self.momentum)
        return success

net = Net()
# 设置优化器待优化的参数和学习率为0.01
opt = Momentum(net.trainable_params(), 0.01)

